"""
ModelSpec = [coef, AlphaSpec, TransformationSpec]

AlphaSpec = [alphapath, params] -> params: prod0, prod1, window_size

TransformationSpec = [transformationpath, params]
-> transformation: zscore_precalculated_sample, params: mu sigma

TransformationType
identity
signed log
signed sqrt
zscore (mu, sigma) parameter
sigmoid (k * x) : k as parameter

*kwargs

Reversion -> 5m vs 60m

Reversion_5m = functools.partial(Reversion, window=5m)

Alpha1 = Reversion, 5m, self
Alpha2 = Reversion, 60m, self
CrossAlpha1 = CrossPrice, (3m, symbola, symbolb), Zscore, (mu, sigma)
    -- mu, sigma = in sample calculated
.
.
.
Alphan ~

[Alpha1 ~ Alphan] & coef, intercept -> model
"""
import numpy as np

from absl import app


class Alpha:
  def __init__(self, func, params):
    self._func = func
    self._params = params

  def eval(self, x):
    return self._func(x, **self._params)

  @property
  def name(self):
    name = f"{self._func.__name__} with {self._params}"
    return name


class Transform:
  def __init__(self, func, params):
    self._func = func
    self._params = params

  def eval(self, x):
    return self._func(x, **self._params)

  @property
  def name(self):
    name = f"{self._func.__name__} with {self._params}"
    return name


class AlphaTrans:
  def __init__(self, alpha, trans=None):
    self._alpha = alpha
    self._trans = trans

  def get_alpha_name(self):
    return self._alpha.__name__

  def eval(self, x):
    return self._trans.eval(self._alpha.eval(x))

  @property
  def alpha(self):
    return self._alpha

  @property
  def transform(self):
    return self._trans


class LinearModel:
  def __init__(self, alpha_trans_lst=[]):
    self._alpha_trans_lst = alpha_trans_lst
    self._coef = []
    self._intercept = 0.0

  def eval(self, row):
    pos = self._intercept

    for coef, alpha_trans in zip(self._coef, self._alpha_trans_lst):
      pos += (alpha_trans.eval(row) * coef)

    return pos

  def add(self, coef, alpha_trans):
    self._coef.append(coef)
    self._alpha_trans_lst.append(alpha_trans)

  def update_intercept(self, intercept):
    self._intercept = intercept

  def update_coef(self, coef):
    self._coef = coef

  @staticmethod
  def ridge_regress_and_predict(X_in, y_in, X_out, y_out):
    from sklearn.linear_model import Ridge
    clf = Ridge(alpha=0.9, normalize=True)
    clf.fit(X_in, y_in)
    y_pred_out = clf.predict(X_out)
    return clf, y_pred_out

  def ridge_fit(self, df, y, is_ratio=0.5):
    x_list = list()
    for alpha_trans in self._alpha_trans_lst:
      x_list.append(alpha_trans.eval(df))
    # X = pd.concat(x_list, axis=1).fillna(0)
    X = np.nan_to_num(np.array(x_list).transpose())
    halfidx = int(len(y) * is_ratio)
    X_in, y_in = X[:halfidx], y[:halfidx]
    X_out, y_out = X[halfidx:], y[halfidx:]
    clf, y_pred_out = LinearModel.ridge_regress_and_predict(X_in, y_in, X_out, y_out)
    y_pred_in = clf.predict(X_in)

    self.update_coef(clf.coef_)
    self.update_intercept(clf.intercept_)

    return np.append(y_pred_in, y_pred_out), [clf.score(X_in, y_in), clf.score(X_out, y_out)]

  def print(self):
    print(self._intercept)
    print(self._coef)
    print(self._alpha_trans_lst)

  def get_coef_alpha_trans(self):
    return sorted(zip(self._coef, self._alpha_trans_lst), key=lambda x: abs(x[0]), reverse=True)

  def save(self, filepath):
    with open(filepath, 'wb') as f_out:
      pickle.dump(self, f_out)

  @staticmethod
  def load(filepath):
    with open(filepath, 'rb') as f:
      model = pickle.load(f)
    return model


from coin.strategy.interval.alphas.taekwon import alpha_1, alpha_2, temp_trans

import pickle
import pandas as pd


def save_test():
  alpha1 = Alpha(alpha_1, {'mul': 2})
  alpha2 = Alpha(alpha_2, {'pow': 3})
  trans1 = Transform(temp_trans, {'offset': 10000})

  model = LinearModel()
  model.add(1.0, AlphaTrans(alpha1, trans1))
  model.add(2.0, AlphaTrans(alpha2, trans1))

  row = pd.Series(range(5))

  print(model.eval(row))
  model.update_intercept(5.5)
  print(model.eval(row))
  model.save('sung.pkl')


def load_test():
  model = LinearModel.load('sung.pkl')
  model.print()
  row = pd.Series(range(5))
  print(model.eval(row))


def main(argv):
  # save_test()
  load_test()


if __name__ == '__main__':
  app.run(main)
